from queue import Queue
from random import random
from generate_benchmark import Task, loadBenchmark, generateBenchmark
from random import random, choice, randint
from config import *

# 默认是随机的assign节点
class Engine():
    def __init__(self, node_number, ):
        self.node_number = node_number  #假设一个node只有一个processor，多个可以当成network speed无限大的
        self.node2queue = [
            Queue()
            for node in range(node_number)
        ]

        self.tasks_sending = []  #正在传输的task  (task, required_time, start_time, target_node)

        self.node2occupied_task = [
            None
            for node in range(node_number)
        ]

        self.now_time = 0

        self.finished_task_num = 0
        return

    def init(self, file_name):
        # generateBenchmark(file_name, 1000)
        root_tasks = loadBenchmark(file_name)
        id2task = root_tasks[0].id2task

        self.id2task = root_tasks[0].id2task
        
        for index, task in enumerate(root_tasks):
            self.sendTask(task, None, index%node_number)
            # self.enqueueTask(task, index%node_number)
        
        self.should_finished_task_number = len(self.id2task)

    def isFinish(self):
        return self.should_finished_task_number <= self.finished_task_num

    def enqueueTask(self, task, node):
        self.node2queue[node].put(task)

    def checkNetwork(self,):
        for send_info in self.tasks_sending:
            task, from_node, to_node = send_info
            if task.finish_send_time <= self.now_time:
                self.tasks_sending.remove(send_info)
                self.enqueueTask(task, to_node)
                print(task, 'from', from_node, 'reachs', to_node)
        return

    def checkExecutingTasks(self):

        required_assigned_tasks = {}

        now_time = self.now_time
        for node, task in enumerate(self.node2occupied_task):
            if task is None:
                queue = self.node2queue[node]
                if not queue.empty():
                    self.execute(queue.get(), node)
                continue

            if task.finish_execute_time == self.now_time:
                self.node2occupied_task[node] = None
                self.finished_task_num += 1
                print(task, 'is finished in', node)
            
            execute_time = now_time - task.start_execute_time
            if execute_time in task.time2son_task:
                son_task = task.time2son_task[execute_time]
                son_task = task.id2task[son_task]
                required_assigned_tasks[node] = son_task
                # .append((son_task, node))

        return required_assigned_tasks


    def step(self):
        print('in step:', self.now_time)

        self.checkNetwork()

        required_assigned_tasks = self.checkExecutingTasks()

        for from_node, task  in required_assigned_tasks.items():
            assign_node =  self.assignNode(task, from_node)
            print(task, 'is called by', task, 'in', from_node, 'assigned to', assign_node)
            self.sendTask(task, from_node, assign_node)

        self.now_time += 1
        return

    def assignNode(self, task, call_node):
        return randint(0, self.node_number-1) #右边是闭的

    def sendTask(self, task, from_node, to_node):
        task.start_send_time = self.now_time
        network_cost = int(task.data_size /  network_speed)
        task.finish_send_time = self.now_time +  network_cost#10是speed

        self.tasks_sending.append((task, from_node, to_node))

        print(task, 'is sended to', to_node, 'cost', network_cost)

    def execute(self, task:Task, node):
        assert 'start_execute_time' not in task.__dict__

        print('execute', task.id, 'in', node)
        task.start_execute_time = self.now_time
        task.finish_execute_time = self.now_time + task.timecost
        task.executor = node

        self.node2occupied_task[node] = task
        return

    def summaryResult(self):
        delays = []
        for task in self.id2task.values():
            delay = task.start_execute_time - task.start_send_time
            delays.append(delay)
        
        print("=======")
        print('averge delay:', sum(delays)/len(delays))

if __name__ == "__main__":
    engine = Engine(node_number)
    engine.init('bechmark100_5.txt')
    # engine.execute(root_tasks[0], 0)
    # for step in range(0, 100):
    #     engine.step()

    while not engine.isFinish():
        engine.step()

    engine.summaryResult()


    


    

    